1092F - Tree with Maximum Cost - CodeForces Solution

from math import inf
from collections import *
import math, os, sys, heapq, bisect, random,threading
from functools import lru_cache
from itertools import *
import sys
def inp(): return sys.stdin.readline().rstrip("\r\n")
def out(var): sys.stdout.write(str(var))  def inpu(): return int(inp())
def lis(): return list(map(int, inp().split()))
def stringlis(): return list(map(str, inp().split()))
def sep(): return map(int, inp().split())
def strsep(): return map(str, inp().split())
def fsep(): return map(float, inp().split())
M,M1=1000000007,998244353


def main():
    how_much_noob_I_am = 1
        for _ in range(1,how_much_noob_I_am+1):
        n = inpu()
        arr = lis()
        d=defaultdict(list)
        for i in range(n-1):
            a,b = sep()
            d[a].append(b)
            d[b].append(a)
        vis = [False]*(n+1)
        ans = [0]*(n+1)
        res = 0
        def dfs(curr,depth):
            nonlocal res
            res+=depth*arr[curr-1]
            ans[curr]=arr[curr-1]
            vis[curr] = True
            for j in d[curr]:
                if not vis[j]:
                    dfs(j,depth+1)
                    ans[curr]+=ans[j]
        dfs(1,0)
        result = 0
        def dfs2(curr,par):
            nonlocal ans,result,res
            result = max(res,result)
            for j in d[curr]:
                if j!=par:
                    res-=ans[j]
                    ans[curr]-=ans[j]
                    res+=ans[curr]
                    ans[j]+=ans[curr]

                    dfs2(j,curr)

                    ans[j] -= ans[curr]
                    res -= ans[curr]
                    ans[curr] += ans[j]
                    res+=ans[j]



        dfs2(1,-1)
        print(result)



if __name__ == '__main__':
    sys.setrecursionlimit(2*10**5+50)
    threading.stack_size(10**8)
    threading.Thread(target=main).start()